// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 *
 * Description: End-to-End Euler Net Boost sample.
 */
#include <linux/bpf.h>
#include <linux/if_link.h>
#include <assert.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <libgen.h>
#include <limits.h>
#include <sys/mount.h>
#include <sys/resource.h>
#include <net/if.h>

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#include "bpf_util.h"
#include <bpf/bpf.h>
#include <bpf/libbpf.h>

#define CGRP_PATH_MODE		0777
#define DEFAULT_IFNAME		"eth0"
#define DEFAULT_BPF_PATH	"bpf.o"
#define DEFAULT_CGRP_PATH	"/sys/fs/cgroup"
#define DEFAULT_SPEED_PORT	5201

struct {
	__u32 ifindex;
	__u32 speed_port;
	char *cgrp_path;
	char *bpf_path;
	bool unload;
	bool help;
} enboost;

struct enboost_prog_info {
	const char *sec_name;
	enum bpf_prog_type prog_type;
	enum bpf_attach_type attach_type;
	int attach_flag;
	int prog_fd;
	bool is_xdp;
};

static struct enboost_prog_info prog_infos[] = {
	{
		.sec_name	= "enboost_sockops",
		.prog_type	= BPF_PROG_TYPE_SOCK_OPS,
		.attach_type	= BPF_CGROUP_SOCK_OPS,
		.attach_flag	= 0,
		.is_xdp		= false,
	},
	{
		.sec_name	= "enboost_ingress",
		.prog_type	= BPF_PROG_TYPE_XDP,
		.attach_type	= BPF_XDP,
		.attach_flag	= XDP_FLAGS_SKB_MODE, // XDP_FLAGS_DRV_MODE or XDP_FLAGS_SKB_MODE
		.is_xdp		= true,
	},
	{
		.sec_name	= "enboost_egress",
		.prog_type	= BPF_PROG_TYPE_CGROUP_EULER_NET,
		.attach_type	= BPF_CGROUP_EULER_NET_EGRESS,
		.attach_flag	= 0,
		.is_xdp		= false,
	},
};

static int set_prog_type(struct bpf_object *obj)
{
	int i;
	const char *sec_name;
	struct bpf_program *prog;
	enum bpf_prog_type prog_type;
	enum bpf_attach_type attach_type;

	bpf_object__for_each_program(prog, obj) {
		sec_name = bpf_program__section_name(prog);
		for (i = 0; i < ARRAY_SIZE(prog_infos); i++) {
			if (!strcmp(prog_infos[i].sec_name, sec_name)) {
				prog_type = prog_infos[i].prog_type;
				attach_type = prog_infos[i].attach_type;
				break;
			}
		}

		if (i == ARRAY_SIZE(prog_infos))
			return -1;

		bpf_program__set_type(prog, prog_type);
		bpf_program__set_expected_attach_type(prog, attach_type);
	}

	return 0;
}

static int find_progs(struct bpf_object *obj)
{
	int i, prog_fd;
	struct bpf_program *prog;
	struct enboost_prog_info *info;

	for (i = 0; i < ARRAY_SIZE(prog_infos); i++) {
		info = &prog_infos[i];
		prog = bpf_object__find_program_by_title(obj, info->sec_name);
		if (!prog) {
			fprintf(stderr, "ERROR: failed to find prog sec %s\n", info->sec_name);
			return -1;
		}

		prog_fd = bpf_program__fd(prog);
		if (prog_fd < 0) {
			fprintf(stderr, "ERROR: failed to get fd of prog %s\n", info->sec_name);
			return -1;
		}

		info->prog_fd = prog_fd;
	}

	return 0;
}

static int set_speed_port(struct bpf_object *obj)
{
	int map_fd;
	__u32 key = 0;
	__u32 val = enboost.speed_port;

	map_fd = bpf_object__find_map_fd_by_name(obj, "speed_port");
	if (map_fd < 0) {
		fprintf(stderr, "ERROR: failed to find map fd\n");
		return -1;
	}

	if (bpf_map_update_elem(map_fd, &key, &val, BPF_ANY)) {
		fprintf(stderr, "ERROR: failed to update map elem\n");
		return -1;
	}

	return 0;
}

static int attach_progs(void)
{
	int i, cgrp_fd;
	struct enboost_prog_info *info;

	cgrp_fd = open(enboost.cgrp_path, O_DIRECTORY, O_RDONLY);
	if (cgrp_fd < 0) {
		fprintf(stderr, "ERROR: failed to open cgrp %s\n", enboost.cgrp_path);
		return -1;
	}

	for (i = 0; i < ARRAY_SIZE(prog_infos); i++) {
		info = &prog_infos[i];
		if (info->is_xdp) {
			if (bpf_set_link_xdp_fd(enboost.ifindex, info->prog_fd, info->attach_flag))
				goto fail;
		} else {
			if (bpf_prog_attach(info->prog_fd, cgrp_fd, info->attach_type,
					    info->attach_flag))
				goto fail;
		}
	}

	close(cgrp_fd);
	return 0;
fail:
	fprintf(stderr, "ERROR: failed to attach prog %s\n", info->sec_name);
	close(cgrp_fd);
	return -1;
}

static int detach_progs(void)
{
	int i, cgrp_fd;
	int err_cnt = 0;
	struct enboost_prog_info *info;

	cgrp_fd = open(enboost.cgrp_path, O_DIRECTORY, O_RDONLY);
	if (cgrp_fd < 0) {
		fprintf(stderr, "ERROR: failed to open cgrp %s\n", enboost.cgrp_path);
		return -1;
	}

	for (i = 0; i < ARRAY_SIZE(prog_infos); i++) {
		info = &prog_infos[i];
		if (info->is_xdp) {
			if (bpf_set_link_xdp_fd(enboost.ifindex, -1, info->attach_flag)) {
				fprintf(stderr,
					"ERROR: failed to detach prog %s\n", info->sec_name);
				err_cnt++;
			}
			continue;
		}

		if (bpf_prog_detach(cgrp_fd, info->attach_type)) {
			fprintf(stderr, "ERROR: failed to detach prog %s\n", info->sec_name);
			err_cnt++;
		}
	}

	close(cgrp_fd);
	return -err_cnt;
}

static int do_enboost(void)
{
	struct bpf_object *obj;

	obj = bpf_object__open(enboost.bpf_path);
	if (libbpf_get_error(obj)) {
		fprintf(stderr, "ERROR: failed to open bpf file\n");
		return -1;
	}

	if (set_prog_type(obj)) {
		fprintf(stderr, "ERROR: failed to set prog type\n");
		bpf_object__close(obj);
		return -1;
	}

	if (bpf_object__load(obj)) {
		fprintf(stderr, "ERROR: failed to load bpf obj\n");
		bpf_object__close(obj);
		return -1;
	}

	if (find_progs(obj)) {
		fprintf(stderr, "ERROR: failed to find progs\n");
		bpf_object__close(obj);
		return -1;
	}

	if (set_speed_port(obj)) {
		fprintf(stderr, "ERROR: failed to set speed port\n");
		bpf_object__close(obj);
		return -1;
	}

	if (attach_progs()) {
		fprintf(stderr, "ERROR: failed to attach progs\n");
		bpf_object__close(obj);
		return -1;
	}

	bpf_object__close(obj);
	return 0;
}

static int setup_cgrp_env(void)
{
	if (mkdir(enboost.cgrp_path, CGRP_PATH_MODE) && errno != EEXIST) {
		fprintf(stderr, "ERROR: failed to mkdir %s\n", enboost.cgrp_path);
		return -1;
	}

	if (mount("none", enboost.cgrp_path, "cgroup2", 0, NULL) && errno != EBUSY) {
		fprintf(stderr, "ERROR: failed to mount cgroup\n");
		return -1;
	}

	return 0;
}

static void do_help(void)
{
	fprintf(stderr,
		"Load:   enboost [-f BPF_FILE] [-c CGRP_PATH] [-p SPEED_PORT] [-i INTERFACE]\n"
		"Unload: enboost -u [-c CGRP_PATH] [-i INTERFACE]\n");
}

static int parse_args(int argc, char **argv)
{
	int opt, port;
	char *ifname = DEFAULT_IFNAME;

	enboost.bpf_path = DEFAULT_BPF_PATH;
	enboost.cgrp_path = DEFAULT_CGRP_PATH;
	enboost.speed_port = DEFAULT_SPEED_PORT;

	while ((opt = getopt(argc, argv, "f:c:p:i:uh")) != -1) {
		switch (opt) {
		case 'f':
			enboost.bpf_path = optarg;
			break;
		case 'c':
			enboost.cgrp_path = optarg;
			break;
		case 'p':
			port = atoi(optarg);
			if (port <= 0 || port >= USHRT_MAX) {
				fprintf(stderr, "ERROR: invalid port: %s\n", optarg);
				return -1;
			}
			enboost.speed_port = port;
			break;
		case 'i':
			ifname = optarg;
			break;
		case 'u':
			enboost.unload = true;
			break;
		case 'h':
			enboost.help = true;
			break;
		default:
			fprintf(stderr, "ERROR: unknown option %c\n", opt);
			return -1;
		}
	}

	enboost.ifindex = if_nametoindex(ifname);
	return 0;
}

int main(int argc, char **argv)
{
	if (parse_args(argc, argv)) {
		fprintf(stderr, "ERROR: failed to parse args\n");
		return -1;
	}

	if (enboost.help) {
		do_help();
		return 0;
	}

	if (setup_cgrp_env()) {
		fprintf(stderr, "ERROR: failed to setup cgroup env\n");
		return -1;
	}

	if (enboost.unload) {
		if (detach_progs()) {
			fprintf(stderr, "ERROR: failed to detach progs\n");
			return -1;
		}

		printf("Unload enboost successfully\n");
		return 0;
	}

	if (do_enboost()) {
		fprintf(stderr, "ERROR: failed to do enboost\n");
		return -1;
	}

	printf("Load enboost successfully\n");
	return 0;
}
